package main

import (
	"bufio"
	"github.com/miekg/dns"
	"log"
	"net"
	"os"
	"strings"
)

// 解析hosts文件并返回域名到IP列表的映射
func parseHostsFile(path string) (map[string]string, error) {
	// 打开文件
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	// 初始化结果map
	hostsMap := make(map[string]string)

	// 逐行扫描文件
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())

		// 跳过空行和注释
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}

		// 分割IP和域名
		fields := strings.Fields(line)
		if len(fields) < 2 {
			continue // 格式不正确，跳过
		}

		ip := fields[0]
		domains := fields[1:]

		// 处理每个域名（移除可能的注释）
		for _, domain := range domains {
			if idx := strings.Index(domain, "#"); idx != -1 {
				domain = strings.TrimSpace(domain[:idx])
			}
			if domain != "" {
				// 添加或更新域名到IP的映射
				hostsMap[domain] = ip
			}
		}
	}

	// 检查扫描错误
	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return hostsMap, nil
}

func getSystemDNSIPList() []string {
	data, err := os.ReadFile("/etc/resolv.conf")
	if err != nil {
		log.Fatal(err)
	}
	lines := strings.Split(string(data), "\n")
	for _, line := range lines {
		if strings.HasPrefix(line, "nameserver") {
			parts := strings.Fields(line)
			return parts[1:]
		}
	}
	return []string{"8.8.8.8"}
}

func startDNSServer() {
	log.Printf("解析hosts文件: %s", hostsFile)
	domainMap, err := parseHostsFile(hostsFile)
	if err != nil {
		log.Fatal(err)
	}
	dns.HandleFunc(".", func(w dns.ResponseWriter, r *dns.Msg) {
		m := new(dns.Msg)
		m.SetReply(r)
		for _, q := range r.Question {
			rawDomain := strings.TrimSuffix(q.Name, ".")
			if ip, ok := domainMap[rawDomain]; ok {
				rr := &dns.A{
					Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 60},
					A:   net.ParseIP(ip), // 自定义IP
				}
				m.Answer = append(m.Answer, rr)
			} else {
				// 转发其他查询到上游DNS
				c := new(dns.Client)
				dnsList := getSystemDNSIPList()
				if len(dnsList) > 0 {
					resp, _, err := c.Exchange(r, dnsList[0]+":53")
					if err != nil {
						log.Fatal(err)
					}
					m.Answer = resp.Answer
				}
			}
		}
		w.WriteMsg(m)
	})
	server := &dns.Server{Addr: dnsServer, Net: "udp"}
	server.ListenAndServe()
}
